/*
 * Copyright (c) Kumo Inc. and affiliates.
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#pragma once

#include <melon/cpp_attributes.h>
#include <melon/function.h>

#include <atomic>
#include <memory>
#include <thread>
#include <type_traits>

namespace melon {
    class CancellationCallback;
    class CancellationSource;

    struct OperationCancelled : public std::exception {
        const char *what() const noexcept override {
            return "coroutine operation cancelled";
        }
    };

    namespace detail {
        class CancellationState;

        struct CancellationStateTokenDeleter {
            void operator()(CancellationState *) noexcept;
        };

        struct CancellationStateSourceDeleter {
            void operator()(CancellationState *) noexcept;
        };

        using CancellationStateTokenPtr =
        std::unique_ptr<CancellationState, CancellationStateTokenDeleter>;
        using CancellationStateSourcePtr =
        std::unique_ptr<CancellationState, CancellationStateSourceDeleter>;
        template<typename...>
        struct WithDataTag;
    } // namespace detail

    /**
     * A CancellationToken is an object that can be passed into an function or
     * operation that allows the caller to later request that the operation be
     * cancelled.
     *
     * A CancellationToken object can be obtained by calling the .getToken()
     * method on a CancellationSource or by copying another CancellationToken
     * object. All CancellationToken objects obtained from the same original
     * CancellationSource object all reference the same underlying cancellation
     * state and will all be cancelled together.
     *
     * If your function needs to be cancellable but does not need to request
     * cancellation then you should take a CancellationToken as a parameter.
     * If your function needs to be able to request cancellation then you
     * should instead take a CancellationSource as a parameter.
     *
     * @refcode melon/docs/examples/melon/CancellationToken.cpp
     * @class melon::CancellationToken
     */
    class CancellationToken {
    public:
        /**
         * Constructs to a token that can never be cancelled.
         *
         * Pass a default-constructed CancellationToken into an operation that
         * you never intend to cancel. These objects are very cheap to create.
         */
        CancellationToken() noexcept = default;

        /// Construct a copy of the token that shares the same underlying state.
        CancellationToken(const CancellationToken &other) noexcept;

        /// Construct a token by moving the underlying state
        CancellationToken(CancellationToken &&other) noexcept;

        CancellationToken &operator=(const CancellationToken &other) noexcept;

        CancellationToken &operator=(CancellationToken &&other) noexcept;

        /**
         * Query whether someone has called .requestCancellation() on an instance
         * of CancellationSource object associated with this CancellationToken.
         */
        bool isCancellationRequested() const noexcept;

        /**
         * Query whether this CancellationToken can ever have cancellation requested
         * on it.
         *
         * This will return false if the CancellationToken is not associated with a
         * CancellationSource object. eg. because the CancellationToken was
         * default-constructed, has been moved-from or because the last
         * CancellationSource object associated with the underlying cancellation state
         * has been destroyed and the operation has not yet been cancelled and so
         * never will be.
         *
         * Implementations of operations may be able to take more efficient code-paths
         * if they know they can never be cancelled.
         */
        bool canBeCancelled() const noexcept;

        /**
         * Obtain a CancellationToken linked to any number of other
         * CancellationTokens.
         *
         * This token will have cancellation requested when any of the passed-in
         * tokens do.
         * This token is cancellable if any of the passed-in tokens are at the time of
         * construction.
         */
        template<typename... Ts>
        static CancellationToken merge(Ts &&... tokens);

        /**
         * Swaps the underlying state of the cancellation token with the token that is
         * passed-in.
         */
        void swap(CancellationToken &other) noexcept;

        friend bool operator==(
            const CancellationToken &a, const CancellationToken &b) noexcept;

    private:
        friend class CancellationCallback;
        friend class CancellationSource;

        explicit CancellationToken(detail::CancellationStateTokenPtr state) noexcept;

        detail::CancellationStateTokenPtr state_;
    };

    bool operator==(
        const CancellationToken &a, const CancellationToken &b) noexcept;

    bool operator!=(
        const CancellationToken &a, const CancellationToken &b) noexcept;

    /**
     * A CancellationSource object provides the ability to request cancellation of
     * operations that an associated CancellationToken was passed to.
     *
     * @refcode melon/docs/examples/melon/CancellationSource.cpp
     * @class melon::CancellationSource
     */
    // Example usage:
    //   CancellationSource cs;
    //   Future<void> f = startSomeOperation(cs.getToken());
    //
    //   // Later...
    //   cs.requestCancellation();
    class CancellationSource {
    public:
        /// Construct to a new, independent cancellation source.
        CancellationSource();

        /**
         * Construct a new reference to the same underlying cancellation state.
         *
         * Either the original or the new copy can be used to request cancellation
         * of associated work.
         */
        CancellationSource(const CancellationSource &other) noexcept;

        /**
         * This leaves 'other' in an empty state where 'requestCancellation()' is a
         * no-op and 'canBeCancelled()' returns false.
         */
        CancellationSource(CancellationSource &&other) noexcept;

        CancellationSource &operator=(const CancellationSource &other) noexcept;

        CancellationSource &operator=(CancellationSource &&other) noexcept;

        /**
         * Construct a CancellationSource that cannot be cancelled.
         *
         * This factory function can be used to obtain a CancellationSource that
         * is equivalent to a moved-from CancellationSource object without needing
         * to allocate any shared-state.
         */
        static CancellationSource invalid() noexcept;

        /**
         * Query if cancellation has already been requested on this CancellationSource
         * or any other CancellationSource object copied from the same original
         * CancellationSource object.
         */
        bool isCancellationRequested() const noexcept;

        /**
         * Query if cancellation can be requested through this CancellationSource
         * object. This will only return false if the CancellationSource object has
         * been moved-from.
         */
        bool canBeCancelled() const noexcept;

        /**
         * Obtain a CancellationToken linked to this CancellationSource.
         *
         * This token can be passed into cancellable operations to allow the caller
         * to later request cancellation of that operation.
         */
        CancellationToken getToken() const noexcept;

        /**
         * Request cancellation of work associated with this CancellationSource.
         *
         * This will ensure subsequent calls to isCancellationRequested() on any
         * CancellationSource or CancellationToken object associated with the same
         * underlying cancellation state to return true.
         *
         * If this is the first call to requestCancellation() on any
         * CancellationSource object with the same underlying state then this call
         * will also execute the callbacks associated with any CancellationCallback
         * objects that were constructed with an associated CancellationToken.
         *
         * Note that it is possible that another thread may be concurrently
         * registering a callback with CancellationCallback. This method guarantees
         * that either this thread will see the callback registration and will
         * ensure that the callback is called, or the CancellationCallback constructor
         * will see the cancellation-requested signal and will execute the callback
         * inline inside the constructor.
         *
         * Returns the previous state of 'isCancellationRequested()'. i.e.
         *  - 'true' if cancellation had previously been requested.
         *  - 'false' if this was the first call to request cancellation.
         */
        bool requestCancellation() const noexcept;

        /**
         * Swaps the underlying state of the cancellation source with the source that
         * is passed-in.
         *
         * @param other The other cancellation source to copy the underlying state
         * from.
         */
        void swap(CancellationSource &other) noexcept;

        friend bool operator==(
            const CancellationSource &a, const CancellationSource &b) noexcept;

        /**
         * Returns a pair of <CancellationSource, Data> where the underlying state is
         * created using the arguments that is passed-in.
         */
        template<typename... Data, typename... Args>
        static std::pair<CancellationSource, std::tuple<Data...> *> create(
            detail::WithDataTag<Data...>, Args &&...);

    private:
        explicit CancellationSource(
            detail::CancellationStateSourcePtr &&state) noexcept;

        detail::CancellationStateSourcePtr state_;
    };

    bool operator==(
        const CancellationSource &a, const CancellationSource &b) noexcept;

    bool operator!=(
        const CancellationSource &a, const CancellationSource &b) noexcept;

    /**
     * A CancellationCallback object registers the callback with the specified
     * CancellationToken such that the callback will be
     * executed if the corresponding CancellationSource object has the
     * requestCancellation() method called on it.
     *
     * If the CancellationToken object already had cancellation requested
     * then the callback will be executed inline on the current thread before
     * the constructor returns. Otherwise, the callback will be executed on
     * in the execution context of the first thread to call requestCancellation()
     * on a corresponding CancellationSource.
     *
     * The callback object must not throw any unhandled exceptions. Doing so
     * will result in the program terminating via std::terminate().
     *
     * A CancellationCallback object is neither copyable nor movable.
     *
     * @refcode melon/docs/examples/melon/CancellationCallback.cpp
     * @class melon::CancellationCallback
     */
    class CancellationCallback {
        using VoidFunction = melon::Function<void()>;

    public:
        template<
            typename Callable,
            std::enable_if_t<
                std::is_constructible<VoidFunction, Callable>::value,
                int>  = 0>
        CancellationCallback(CancellationToken &&ct, Callable &&callable);

        template<
            typename Callable,
            std::enable_if_t<
                std::is_constructible<VoidFunction, Callable>::value,
                int>  = 0>
        CancellationCallback(const CancellationToken &ct, Callable &&callable);

        /**
         * Deregisters the callback from the CancellationToken.
         *
         * If cancellation has been requested concurrently on another thread and the
         * callback is currently executing then the destructor will block until after
         * the callback has returned (otherwise it might be left with a dangling
         * reference).
         *
         * You should generally try to implement your callback functions to be lock
         * free to avoid deadlocks between the callback executing and the
         * CancellationCallback destructor trying to deregister the callback.
         *
         * If the callback has not started executing yet then the callback will be
         * deregistered from the CancellationToken before the destructor completes.
         *
         * Once the destructor returns you can be guaranteed that the callback will
         * not be called by a subsequent call to 'requestCancellation()' on a
         * CancellationSource associated with the CancellationToken passed to the
         * constructor.
         */
        ~CancellationCallback();

        // Not copyable/movable
        CancellationCallback(const CancellationCallback &) = delete;

        CancellationCallback(CancellationCallback &&) = delete;

        CancellationCallback &operator=(const CancellationCallback &) = delete;

        CancellationCallback &operator=(CancellationCallback &&) = delete;

    private:
        friend class detail::CancellationState;

        void invokeCallback() noexcept;

        CancellationCallback *next_;

        // Pointer to the pointer that points to this node in the linked list.
        // This could be the 'next_' of a previous CancellationCallback or could
        // be the 'head_' pointer of the CancellationState.
        // If this node is inserted in the list then this will be non-null.
        CancellationCallback **prevNext_;

        detail::CancellationState *state_;
        VoidFunction callback_;

        // Pointer to a flag stored on the stack of the caller to invokeCallback()
        // that is used to indicate to the caller of invokeCallback() that the
        // destructor has run and it is no longer valid to access the callback
        // object.
        bool *destructorHasRunInsideCallback_;

        // Flag used to signal that the callback has completed executing on another
        // thread and it is now safe to exit the destructor.
        std::atomic<bool> callbackCompleted_;
    };
} // namespace melon

#include <melon/cancellation_token-inl.h>
